import numpy as np
import torch

from torchvision.transforms import Compose

from utils.utils import Rescale, Normailize, Reshape
# from utils.nms_wrapper import nms
   
import torch.nn.functional as F
import torch
import os
import numpy as np
import math
from datasets.DOTA_devkit.ResultMerge_multi_process import py_cpu_nms_poly_fast, py_cpu_nms_poly
from PIL import Image, ImageDraw
import cv2

def decode_prediction(predictions):

    for class_name,lx,ly,rx,ry,ang, prob in predictions:
        result = [int((rx+lx)/2),int((ry+ly)/2),int(rx-lx),int(ry-ly),ang]
        result=np.array(result)
        x=int(result[0])
        y=int(result[1])
        height=int(result[3])
        width=int(result[2])
        anglePi = result[4]/180 * math.pi
        anglePi = anglePi if anglePi <= math.pi else anglePi - math.pi         
        cosA = math.cos(anglePi)
        sinA = math.sin(anglePi)
        
        x1=x-0.5*width   
        y1=y-0.5*height
        
        x0=x+0.5*width 
        y0=y1
        
        x2=x1            
        y2=y+0.5*height 
        
        x3=x0   
        y3=y2
        
        x0n= (x0 -x)*cosA -(y0 - y)*sinA + x
        y0n = (x0-x)*sinA + (y0 - y)*cosA + y        
        x1n= (x1 -x)*cosA -(y1 - y)*sinA + x
        y1n = (x1-x)*sinA + (y1 - y)*cosA + y        
        x2n= (x2 -x)*cosA -(y2 - y)*sinA + x
        y2n = (x2-x)*sinA + (y2 - y)*cosA + y        
        x3n= (x3 -x)*cosA -(y3 - y)*sinA + x
        y3n = (x3-x)*sinA + (y3 - y)*cosA + y     
        
        tr = np.asarray([x0n,y0n], np.float32)
        br = np.asarray([x3n,y3n], np.float32)
        bl = np.asarray([x2n,y2n], np.float32)
        tl = np.asarray([x1n,y1n], np.float32)
        cen_pt = np.asarray([x, y], np.float32)
        
        score = prob            
        clse = int(class_name)
#         print(clse)
        pts = [clse,score,x1n,y1n,x0n,y0n,x3n,y3n,x2n,y2n]      
    return pts

def get_3rd_point(a, b):
    direct = a - b
    return b + np.array([-direct[1], direct[0]], dtype=np.float32)

def get_dir(src_point, rot_rad):
    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
    src_result = [0, 0]
    src_result[0] = src_point[0] * cs - src_point[1] * sn
    src_result[1] = src_point[0] * sn + src_point[1] * cs
    return src_result    
    
def get_affine_transform(center,scale,rot, output_size,
                         shift=np.array([0, 0], dtype=np.float32),inv=0):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        scale = np.array([scale, scale], dtype=np.float32)
    scale_tmp = scale
    src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]
    rot_rad = np.pi * rot / 180 #=0
    src_dir = get_dir([0, src_w * -0.5], rot_rad)#乘以-0.5
    dst_dir = np.array([0, dst_w * -0.5], np.float32)
    
    src = np.zeros((3, 2), dtype=np.float32)#初始化矩阵。shape（3，2）
    dst = np.zeros((3, 2), dtype=np.float32)
    
    src[0, :] = center + scale_tmp * shift
    src[1, :] = center + src_dir + scale_tmp * shift
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir
    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
    if inv:
        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
    return trans

def affine_transform(pt, t):
    new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T
    new_pt = np.dot(t, new_pt)
    return new_pt[:2]

def transform_preds(coords, center, scale, output_size):
    target_coords = np.zeros(coords.shape)
    trans = get_affine_transform(center, scale, 0, output_size, inv=1)
    trans[:,-1]=0. 
    for p in range(coords.shape[0]):
        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
    return target_coords

def ctdet_post_process(dets, c, s, h, w, num_classes):

    ret = []
    for i in range(dets.shape[0]):
        top_preds = {}

        dets[i, :, :2] = transform_preds(dets[i, :, 0:2], c[i], s[i], (w, h))   
        dets[i, :, 2:4] = transform_preds(dets[i, :, 2:4], c[i], s[i], (w, h))
        classes = dets[i, :, -1]

    for j in range(num_classes):
        inds = (classes == j)
        top_preds[j + 1] = np.concatenate([
        dets[i, inds, :4].astype(np.float32),
        dets[i, inds, 4:6].astype(np.float32)], axis=1).tolist()
        ret.append(top_preds)

    return ret


def post_process(dets, meta,num_classes):  

    dets = dets.reshape(1, -1, dets.shape[2])     
    dets = ctdet_post_process(dets.copy(), [meta['c']], [meta['s']],meta['out_height'], meta['out_width'], num_classes)
    for j in range(1, num_classes + 1):
        dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 6)
        dets[0][j][:, :5] /= 1

    return dets[0]

def merge_outputs(detections,num_classes):

    max_obj_per_img = 100
    scores = np.hstack([detections[j][:, 5] for j in range(1, num_classes + 1)])#在水平方向上平铺
    if len(scores) > max_obj_per_img:
        kth = len(scores) - max_obj_per_img
        thresh = np.partition(scores, kth)[kth]
        for j in range(1, 2 + 1):
            keep_inds = (detections[j][:, 5] >= thresh)
            detections[j] = detections[j][keep_inds]
    return detections


def processing_test(image, inp_height, inp_width):
        
    height, width = image.shape[0:2]
    c = np.array([height / 2., width / 2.], dtype=np.float32)
    s = max(height, width) * 1     
    trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])    
    trans_input[:, -1] = 0.
    inp_image = cv2.warpAffine(image, trans_input, (inp_width, inp_height), flags=cv2.INTER_LINEAR)

    mean = np.array([0.5194416012442385, 0.5378052387430711, 0.533462090585746], dtype=np.float32).reshape(1, 1, 3)
    std = np.array([0.3001546018824507, 0.28620901391179554, 0.3014112676161966], dtype=np.float32).reshape(1, 1, 3)

    inp_image = ((inp_image / 255. - mean) / std).astype(np.float32)

    images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)  # 三维reshape到4维，（1，3，512，512） 

    images = torch.from_numpy(images)
        
    meta = {'c': c, 's': s,
                'out_height': inp_height // 4,  # 取整除
                'out_width': inp_width // 4}

    return images,meta


def im_detect(args,
              image,
              model,
              down_ratio):
  
    decoder = DecDecoder(100,0.1,args.num_classes)  #K, conf_thresh, num_classes
    
    image,meta = processing_test(image, args.input_h, args.input_w)
    if torch.cuda.is_available():
        model, image = model.cuda(), image.cuda()
        
    with torch.no_grad():
        pr_decs = model(image)#pr_decs[ang].shape = (1,180) model:angle:180
    pr_decs['hm'] = pr_decs['hm'].sigmoid_()      
    if args.CSL:       
        pr_decs['ang'] = pr_decs['ang'].sigmoid_()        
    else:
        pr_decs['ang'] = pr_decs['ang'].relu_()

    decoded_pts = []
    decoded_scores = []
    
        
    predictions = decoder.ctdet_decode(pr_decs,args)
    predictions = post_process(predictions, meta,args.num_classes)
    
    ret = merge_outputs(predictions,args.num_classes)#没用 但是最好不删
  
    bboxes = []
    for i, c in ret.items():
        tmp_s = ret[i][ret[i][:,5]>0.15]
        res = np.empty([1,7])
        tmp_c = np.ones(len(tmp_s)) * (i)
       
        tmp = np.c_[tmp_c,tmp_s]
        res = np.append(res,tmp,axis=0)
        res = np.delete(res, 0, 0)
        res = res.tolist()
        
        for bbox in res:          
            pts = decode_prediction([bbox])
            bboxes.append(pts)

     
    bboxes = np.asarray(bboxes, np.float32)
    return bboxes

class DecDecoder(object):
    def __init__(self, K, conf_thresh, num_classes):
        self.K = K
        self.conf_thresh = conf_thresh
        self.num_classes = num_classes

    def _topk(self, scores):
        batch, cat, height, width = scores.size()

        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), self.K)

        topk_inds = topk_inds % (height * width)
        topk_ys = (topk_inds // width).int().float()
        topk_xs = (topk_inds % width).int().float()

        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), self.K)
        topk_clses = (topk_ind // self.K).int()
        topk_inds = self._gather_feat( topk_inds.view(batch, -1, 1), topk_ind).view(batch, self.K)
        topk_ys = self._gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, self.K)
        topk_xs = self._gather_feat(topk_xs.view  (batch, -1, 1), topk_ind).view(batch, self.K)

        return topk_score, topk_inds, topk_clses, topk_ys, topk_xs


    
    def _nms(self, heat, kernel=3):
        hmax = F.max_pool2d(heat, (kernel, kernel), stride=1, padding=(kernel - 1) // 2)
        keep = (hmax == heat).float()
        return heat * keep
    


    def _gather_feat(self, feat, ind, mask=None):
        dim = feat.size(2)
        ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
        feat = feat.gather(1, ind)
        if mask is not None:
            mask = mask.unsqueeze(2).expand_as(feat)
            feat = feat[mask]
            feat = feat.view(-1, dim)
        return feat

    def _transpose_and_gather_feat(self, feat, ind):
        feat = feat.permute(0, 2, 3, 1).contiguous()
        feat = feat.view(feat.size(0), -1, feat.size(3))
        feat = self._gather_feat(feat, ind)
        return feat

    def ctdet_decode(self, pr_decs,args):

        heat = pr_decs['hm']
        wh = pr_decs['wh']
        reg = pr_decs['reg']
        ang = pr_decs['ang']

        batch, cat, height, width = heat.size() # heat = torch.sigmoid(heat) # perform nms on heatmaps
        heat = self._nms(heat)#yiyang
        scores, inds, clses, ys, xs = self._topk(heat)
        reg = self._transpose_and_gather_feat(reg, inds)
        reg = reg.view(batch, self.K, 2)
        xs = xs.view(batch, self.K, 1) + reg[:, :, 0:1]
        ys = ys.view(batch, self.K, 1) + reg[:, :, 1:2]
        wh =  self._transpose_and_gather_feat(wh, inds)
        wh = wh.view(batch, self.K, 2)
        ang =  self._transpose_and_gather_feat(ang, inds)
        if args.CSL:
            ang = ang.view(batch, self.K, 180)

            ang =torch.argmax(ang,dim=2).view(batch, self.K, 1)
            ang = ang.float()

        else:
            ang = ang.view(batch, self.K, 1)

        clses  = clses.view(batch, self.K, 1).float()
        scores = scores.view(batch, self.K, 1)
        bboxes = torch.cat([xs - wh[..., 0:1] / 2, 
                            ys - wh[..., 1:2] / 2,
                            xs + wh[..., 0:1] / 2, 
                            ys + wh[..., 1:2] / 2,
                            ang], dim=2)
        detections = torch.cat([bboxes, scores, clses], dim=2)
        return detections.detach().cpu().numpy()
    
    
    
    








